import os
import tensorflow as tf
import tensorflow.contrib.slim as slim

DATASET_DIR = '/home/raintai/guoyf/VOC/VOCdevkit/VOC2012'
FILE_PATTERN = 'voc_%s_*.tfrecord'
ITEMS_TO_DESCRIPTIONS = {
    'image': 'A color image of varying height and width.',
    'shape': 'Shape of the image',
    'object/bbox': 'A list of bounding boxes, one per each object.',
    'object/label': 'A list of labels, one per each object.',
}
NUM_CLASSES = 20
SPLIT_NAME = 'train' # 'train' or 'test'

"""
Gets a dataset tuple with instructions for reading Pascal VOC dataset.

Args:
  split_name:   A train/test split name.
  dataset_dir:  The base directory of the dataset sources.
  file_pattern: The file pattern to use when matching the dataset sources.
                It is assumed that the pattern contains a '%s' string so that the split
                name can be inserted.
  reader:       The TensorFlow reader type.

Returns:
  A `Dataset` namedtuple.

Raises:
    ValueError: if `split_name` is not a valid train/test split.
"""
def get_split():
    # Features in Pascal VOC TFRecords.
    keys_to_features = {
        'image/encoded': tf.FixedLenFeature((), tf.string, default_value=''),
        'image/format': tf.FixedLenFeature((), tf.string, default_value='jpeg'),
        'image/height': tf.FixedLenFeature([1], tf.int64),
        'image/width': tf.FixedLenFeature([1], tf.int64),
        'image/channels': tf.FixedLenFeature([1], tf.int64),
        'image/shape': tf.FixedLenFeature([3], tf.int64),
        'image/object/bbox/xmin': tf.VarLenFeature(dtype=tf.float32),
        'image/object/bbox/ymin': tf.VarLenFeature(dtype=tf.float32),
        'image/object/bbox/xmax': tf.VarLenFeature(dtype=tf.float32),
        'image/object/bbox/ymax': tf.VarLenFeature(dtype=tf.float32),
        'image/object/bbox/label': tf.VarLenFeature(dtype=tf.int64),
        'image/object/bbox/difficult': tf.VarLenFeature(dtype=tf.int64),
        'image/object/bbox/truncated': tf.VarLenFeature(dtype=tf.int64),
    }
    items_to_handlers = {
        'image': slim.tfexample_decoder.Image('image/encoded', 'image/format'),
        'shape': slim.tfexample_decoder.Tensor('image/shape'),
        'object/bbox': slim.tfexample_decoder.BoundingBox(
                ['ymin', 'xmin', 'ymax', 'xmax'], 'image/object/bbox/'),
        'object/label': slim.tfexample_decoder.Tensor('image/object/bbox/label'),
        'object/difficult': slim.tfexample_decoder.Tensor('image/object/bbox/difficult'),
        'object/truncated': slim.tfexample_decoder.Tensor('image/object/bbox/truncated'),
    }
    decoder = slim.tfexample_decoder.TFExampleDecoder(keys_to_features, items_to_handlers)


    return slim.dataset.Dataset(
    			data_sources = os.path.join(DATASET_DIR, FILE_PATTERN % SPLIT_NAME),
				reader = tf.TFRecordReader,
				decoder=decoder,
				num_samples = 17125,
				items_to_descriptions = ITEMS_TO_DESCRIPTIONS,
				num_classes = NUM_CLASSES,
				labels_to_names=None
				)

if __name__ == "__main__":
	dataset = get_split();
	print 'OK'









